#!/usr/bin/env python
from astropy.io import fits
from glob import glob
import os,sys
import shutil
import numpy as np
from astropy.io import fits
from scipy.signal import savgol_filter
from scipy.stats import trim_mean, sigmaclip
from tqdm import trange

class ConfGM1201A:
    def __init__(self):
        self.name = 'gm1201a'
        self.type = 'ccd'
        self.tsys = 188
        self.glass = 'none'
        self.wave = -99.0

        # channel layout
        self.nchan = 16
        self.col = [0, 1, 2, 3, 4, 5, 6, 7, 7, 6, 5, 4, 3, 2, 1, 0]
        self.row = [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]
        self.overscan = 1

        # raw image
        self.raw_shape = (9400, 10000)
        self.raw_x0 = [1250 * i for i in self.col]
        self.raw_y0 = [4700 * i for i in self.row]
        self.raw_nx = [1250] * 16
        self.raw_ny = [4700] * 16
        self.raw_xflip = [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0]  # flip for overscan region alignment
        self.raw_yflip = [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]

        # exposure region, prescan and overscan (in raw-flipped channel cutout)
        self.x1 = [29] * 16
        self.y1 = [4] * 16
        self.nx = [1152] * 16
        self.ny = [4616] * 16
        self.overscan_x1 = [1200] * 16
        self.overscan_x2 = [1250] * 16
        self.overscan_y1 = [4650] * 16
        self.overscan_y2 = [4700] * 16

        # mosaic image (without prescan and overscan)
        self.mos_shape = (9232, 9216)
        self.mos_x0 = [1152 * i for i in self.col]
        self.mos_y0 = [4616 * i for i in self.row]
        self.mos_xflip = [0] * 16  # flip the exposure region for mosaic
        self.mos_yflip = [0] * 16
        
        # image without prescan
        self.nop_shape = (9392,9768)
        self.nop_x0 = [1221 * i for i in self.col] 
        self.nop_y0 = [4696 * i for i in self.row]
        self.nop_nx = [1221] * 16
        self.nop_ny = [4696] * 16
        
        # image without prescan and parallel overscan
        self.nopp_shape = (9232,9768)
        self.nopp_x0 = [1221 * i for i in self.col] 
        self.nopp_y0 = [4616 * i for i in self.row]
        self.nopp_nx = [1221] * 16
        self.nopp_ny = [4616] * 16
        
class ConfHR9090BSI:
    def __init__(self):
        self.name = 'hr9090bsi'
        self.type = 'cmos'
        self.tsys = 188
        self.glass = 'none'
        self.wave = -99.0

        # channel layout
        self.nchan = 2
        self.col = [0, 1]
        self.row = [0, 0]
        self.overscan = 1

        # raw image
        self.raw_shape = (9222, 8976)
        self.raw_x0 = [0, 4352]
        self.raw_y0 = [0, 0]
        self.raw_nx = [4352, 4624]
        self.raw_ny = [9222, 9222]
        self.raw_xflip = [0, 1]  # flip for overscan region alignment
        self.raw_yflip = [0, 0]

        # exposure region, prescan and overscan (in raw-flipped channel cutout)
        self.x1 = [38, 38]  # EB 34 + dummy 4
        self.y1 = [92, 92]  # dummy 92
        self.nx = [4314, 4586]
        self.ny = [9120, 9120]
        self.overscan_x1 = [0, 0]
        self.overscan_x2 = [34, 34]  # only use EB
        self.overscan_y1 = [0, 0]
        self.overscan_y2 = [92, 92]

        # mosaic image (without prescan and overscan)
        self.mos_shape = (9120, 8900)
        self.mos_x0 = [0, 4314]
        self.mos_y0 = [0, 0]
        self.mos_xflip = [0, 0]  # flip the exposure region for mosaic
        self.mos_yflip = [0, 0]
       
        # image without prescan
        self.nop_shape = (9212, 8968)
        self.nop_x0 = [0, 4348]
        self.nop_y0 = [0, 0]
        self.nop_nx = [4348, 4620]
        self.nop_ny = [9212, 9212]

        # image without prescan
        self.nopp_shape = (9154,8832)
        self.nopp_x0 = [0, 4348] 
        self.nopp_y0 = [0, 0]   
        self.nopp_nx = [4348, 4484]
        self.nopp_ny = [9120, 9120]    
        
def _smooth_data(data, filterwidth):
    dd = savgol_filter(data, filterwidth * 4 + 1, 3)
    diff = data - dd
    index = np.abs(diff) > sigmaclip(diff, 5, 5)[0].std() * 5
    x = np.arange(len(data))
    dd = np.interp(x, x[~index], data[~index])
    return savgol_filter(dd, filterwidth * 2 + 1, 2)


def extract_overscan(image, os_x1, os_x2, y1, y2, filterwidth=5):
    if os_x1 >= os_x2:
        raise Exception('os_x1 must be smaller than os_x2')
    if y1 >= y2:
        raise Exception('y1 must be smaller than y2')
    data = trim_mean(image[y1:y2, os_x1:os_x2], 0.1, axis=1)
    if filterwidth > 1:
        data = _smooth_data(data, filterwidth)
    return data

def subtract_overscan_2(image, raw_x0, raw_nx, raw_y0, raw_ny, os_x1, os_x2):
    v = extract_overscan(image, os_x1, os_x2, raw_y0, raw_y0 + raw_ny)
    v = np.repeat(v, raw_nx).reshape((raw_ny, raw_nx))
    image0 = image.copy()
    image0[raw_y0:raw_y0 + raw_ny, raw_x0:raw_x0+raw_nx] -= v
    return image0

def corr_overscan_2(image, imgconf):
    for chan in range(imgconf.nchan):
        x0 = imgconf.raw_x0[chan]
        y0 = imgconf.raw_y0[chan]
        if imgconf.raw_xflip[chan]:
            os1 = x0 + imgconf.raw_nx[chan] - imgconf.overscan_x2[chan]
            os2 = x0 + imgconf.raw_nx[chan] - imgconf.overscan_x1[chan]
        else:
            os1 = x0 + imgconf.overscan_x1[chan]
            os2 = x0 + imgconf.overscan_x2[chan]
        image = subtract_overscan_2(image, imgconf.raw_x0[chan], imgconf.raw_nx[chan], imgconf.raw_y0[chan], imgconf.raw_ny[chan], os1, os2)    
    return image

def crop_prescan(image, imgconf):
    output = np.zeros(imgconf.nop_shape)
    for chan in range(imgconf.nchan):
        x0 = imgconf.raw_x0[chan]
        y0 = imgconf.raw_y0[chan]
        x0_out = imgconf.nop_x0[chan]
        x1_out = imgconf.nop_x0[chan] + imgconf.nop_nx[chan]
        y0_out = imgconf.nop_y0[chan]
        y1_out = imgconf.nop_y0[chan] + imgconf.nop_ny[chan]
        nop_nx = imgconf.nop_nx[chan]
        nop_ny = imgconf.nop_ny[chan]
        if imgconf.raw_xflip[chan]:
            x1 = x0
        else:
            x1 = x0 + imgconf.x1[chan]
        if imgconf.raw_yflip[chan]:
            y1 = y0
        else: 
            y1 = y0 + imgconf.y1[chan]
        cimage = image[y1:y1+nop_ny,x1:x1+nop_nx]
        print("channel",chan+1,"cimage_shape",cimage.shape)
        output[y0_out:y1_out,x0_out:x1_out] = cimage
    return output  

def crop_dummy(image):
    output = np.zeros((9212,8968))
    output[:,:34] = image[:-10,:34]
    output[:,34:-34] = image[:-10,38:-38]
    output[:,-34:] = image[:-10,-34:]
    return output

def crop_overscan(image, imgconf):
      output = np.zeros(imgconf.mos_shape, dtype=image.dtype)
      for chan in range(imgconf.nchan):
          x1 = imgconf.raw_x0[chan]
          y1 = imgconf.raw_y0[chan]
          cimage = image[y1:y1 + imgconf.raw_ny[chan], x1:x1 + imgconf.raw_nx[chan]] 
          if imgconf.raw_xflip[chan]:
              x1 = imgconf.raw_nx[chan] - imgconf.x1[chan] - imgconf.nx[chan]
          else:
              x1 = imgconf.x1[chan]
          if imgconf.raw_yflip[chan]:
              y1 = imgconf.raw_ny[chan] - imgconf.y1[chan] - imgconf.ny[chan]
          else:
              y1 = imgconf.y1[chan]
          cimage = cimage[y1:y1 + imgconf.ny[chan], x1:x1 + imgconf.nx[chan]]
  
          if imgconf.mos_xflip[chan]:
              cimage = cimage[:, ::-1]
          if imgconf.mos_yflip[chan]:
              cimage = cimage[::-1, :]
          x1 = imgconf.mos_x0[chan]
          y1 = imgconf.mos_y0[chan]
          output[y1:y1 + imgconf.ny[chan], x1:x1 + imgconf.nx[chan]] = cimage
      return output

def crop_prescan_and_poverscan(image, imgconf):
    output = np.zeros(imgconf.nopp_shape)
    for chan in range(imgconf.nchan):
        x0 = imgconf.raw_x0[chan]
        y0 = imgconf.raw_y0[chan]
        x0_out = imgconf.nopp_x0[chan]
        x1_out = imgconf.nopp_x0[chan] + imgconf.nopp_nx[chan]
        y0_out = imgconf.nopp_y0[chan]
        y1_out = imgconf.nopp_y0[chan] + imgconf.nopp_ny[chan]
        nopp_nx = imgconf.nopp_nx[chan]
        ny = imgconf.ny[chan]
        if imgconf.raw_xflip[chan]:
            x1 = x0
        else:
            x1 = x0 + imgconf.x1[chan]
        if imgconf.raw_yflip[chan]:
            y1 = y0 + imgconf.raw_ny[chan] - imgconf.y1[chan] - imgconf.ny[chan]
        else: 
            y1 = y0 + imgconf.y1[chan]
        cimage = image[y1:y1+ny,x1:x1+nopp_nx]
        print("channel",chan+1,"cimage_shape",cimage.shape)
        output[y0_out:y1_out,x0_out:x1_out] = cimage
    return output

if __name__ == '__main__':
    conf = sys.argv[1]
    if os.path.exists(sys.argv[2]):
        os.system("rm "+sys.argv[2])
    files = glob("*.fits")
    if conf == '44s':
        imgconf = ConfGM1201A()
    if conf == 'cmos':
        imgconf = ConfHR9090BSI()
    for i in range(len(files)):
        image = fits.getdata(files[i]).astype(float)
        if conf == '44s':
            image_corr = corr_overscan_2(image,imgconf)
        if conf == 'cmos':
            image_corr = image
        if i==0:
            image_stack = image_corr.copy()
        else:
            image_stack += image_corr.copy()
    image_stack /= len(files)
    if conf == '44s':
        image_stack_nopre = crop_prescan(image_stack, imgconf)
    if conf == 'cmos':
        image_stack_nopre = crop_overscan(image, imgconf)
    fits.writeto(sys.argv[2],data=image_stack,overwrite=True)
